Summary

This PR adds SM120 fragment-MMA GEMM support for T.float4_e2m1fn.
It covers plain FP4 GEMM and explicit mixed FP8/FP4 GEMM while keeping
the TileLang-facing API dtype-semantic.

Kernels continue to declare FP4 operands as T.float4_e2m1fn. Packed
byte storage is handled by lowering/codegen and by host-side example setup;
users do not have to model FP4 GEMM operands as uint8 tensors in TileLang
programs.

Supported GEMM combinations:

Design Goals

• Preserve semantic TileLang signatures: FP4 tensors stay T.float4_e2m1fn
  at the language level.
• Keep packed FP4 as a storage/lowering concern, not a public GEMM dtype
  workaround.
• Add the SM120-specific hardware path without changing existing int4/uint4
  ldmatrix behavior.
• Make the implementation reviewable by focusing this PR on SM120 FP4/A8W4
  GEMM.
• Fail early for unsupported FP4/A8W4 MMA tile shapes instead of generating
  kernels that silently skip data.

Hardware Contracts

The SM120 FP4 path has three contracts that must line up across lowering,
layout, and template dispatch:

The key implementation detail is that SM120 b4x16_p64 consumes packed FP4
bytes from shared memory with a padded shared row layout. Global memory,
shared memory, and local fragments therefore cannot all use the same offset
model:

• Global FP4 storage is packed by logical nibbles.
• Shared FP4 storage uses the padded layout required by SM120 ldmatrix.
• Local fragments keep the declared semantic names and types expected by
  the MMA lowering path.

Implementation Details

Semantic FP4 Storage Model

This PR separates "what the TileLang program declares" from "how the FP4
payload is physically moved":

• Language-level operands use T.float4_e2m1fn.
• Host examples pack FP4 values into byte storage only as an interoperability
  detail.
• Codegen recognizes packed global FP4 storage when computing memory offsets.
• Shared-memory copies use the SM120 padded FP4 layout.
• Local fragments are not rewritten into user-visible _packed aliases.

This avoids leaking uint8 into the public GEMM dtype model while still
allowing the generated CUDA path to use the byte/nibble representation that
the hardware requires.

SM120 FP4 LDSM

src/tl_templates/cuda/ldsm.h adds SM120 ptx_ldmatrix_b4x16_x{1,2,4}
helpers guarded for the target architecture. CUDA lowering selects these
helpers when the source fragment is explicitly float4_e2m1fn.

The Python layout side adds FP4-specific ldmatrix logical layouts and offset
handling. This is gated on float4_e2m1fn, so the existing int4/uint4
ldmatrix offset behavior stays on the previous path.

SM120 MMA Dispatch

The template dispatch adds SM120 cute::SM120_16x8x32_TN support for:

• FP4 x FP4 -> FP32
• FP8 e4m3 x FP4 -> FP32
• FP4 x FP8 e4m3 -> FP32

The FP4 operand register payload adjustment is applied only to FP4 operands
before calling the CuTe atom. Mixed A8W4/W4A8 dispatch is selected from the
explicit A/B dtype pair, rather than inferred from a packed integer carrier.

Copy And Async Copy

The copy path distinguishes packed global FP4 offsets from padded shared FP4
offsets. Shared-to-fragment copy lowering routes SM120 FP4 through the new
b4x16_p64 ldmatrix path, while global-to-shared async copy keeps using the
existing cp.async lowering with FP4 padded-shared metadata enabled.

For FP4 global-to-shared async copy, the lowering emits 8-byte segments for
packed FP4 storage and carries the extra metadata needed to place those bytes
into the padded shared-memory layout.

K Tile Validation

SM120 FP4/A8W4 MMA consumes K in instruction-sized m16n8k32 tiles. A
T.gemm block K that is not divisible by the selected instruction K tile
cannot be represented by the current lowering without dropping the leftover
K range.

This PR therefore rejects unsupported K tile choices up front. For example,
block_K=48 is invalid for this path because it would execute one K=32 MMA
tile and silently omit the remaining K=16 tail.

Main Changes

CUDA Templates

CUDA Lowering

Python Lowering

Examples

Changes Added After Initial Review

Why The Review Fixes Matter

FP4 global packed storage is byte-addressed, but logical FP4 elements are
nibble-addressed. A vector reinterpret load/store is only safe when the
logical base offset is known to be even. If the offset is odd, or if codegen
cannot prove that it is even, vectorized byte reinterpretation can read or
write the wrong nibble without producing a compilation error.

Likewise, SM120 FP4/A8W4 MMA consumes K in fixed instruction-sized chunks.
Allowing a block_K such as 48 would make the generated kernel execute only
the representable K=32 portion and miss the K tail. The new validation turns
that silent numerical error into an explicit unsupported-shape error.

Validation

Local SM120 validation used an RTX 5090 / compute capability 12.0 environment.

Build and examples:

cmake --build build -j$(nproc)
PYTHONPATH=$PWD${PYTHONPATH:+:$PYTHONPATH} python examples/gemm_fp4/example_gemm_fp4_sm120.py
PYTHONPATH=$PWD${PYTHONPATH:+:$PYTHONPATH} python examples/gemm_fp4/example_gemm_a8w4_sm120.py

Observed numerical results:

example_gemm_fp4_sm120.py: max_abs_diff=0.000000
example_gemm_a8w4_sm120.py: max_abs_diff=0.000000, rel_err=0.000000

Generated CUDA was inspected for the expected SM120 FP4 markers:

tl::mma_sync<..., 16, 8, 32, false, true>
ptx_ldmatrix_b4x16_x4
FP4xFP4 and FP8xFP4 dtype dispatch

Focused tests:

PYTHONPATH=$PWD${PYTHONPATH:+:$PYTHONPATH} python -m pytest testing/python/kernel/test_tilelang_kernel_fp4_gemm.py -q
PYTHONPATH=$PWD${PYTHONPATH:+:$PYTHONPATH} python -m pytest testing/python/language/test_tilelang_language_copy.py -q -k 'fp4_odd_vector_start or copy_fp4'

Observed focused-test results:

testing/python/kernel/test_tilelang_kernel_fp4_gemm.py: 2 passed
testing/python/language/test_tilelang_language_copy.py -k 'fp4_odd_vector_start or copy_fp4': 2 passed

Notes And Non-Goals

• Mixed A8W4/W4A8 dispatch is selected from explicit FP8/FP4 dtype pairs.
• FP4 ldmatrix layout handling is gated on float4_e2m1fn.
• Existing int4/uint4 ldmatrix offset behavior stays on the existing path.

Status

All set now！